import copy
from lxml import etree, objectify
import cv2
import os
from shutil import copy


# 修改为你自己的路径
# template_file = 'C:\\Users\\zzw\\Desktop\\jongjv\\txt-xml\\000.xml'
# target_dir = 'C:\\Users\\zzw\\Desktop\\jongjv\\txt-xml\\Annotations\\'
# image_dir = 'C:\\Users\\zzw\\Desktop\\jongjv\\txt-xml\\images\\'  # 图片文件夹
# train_file = 'C:\\Users\\zzw\\Desktop\\jongjv\\txt-xml\\123.txt'  # 存储了图片信息的txt文件

def txt2xml(template_file, file_name, filepath, train_file):
    with open(train_file) as f:
        trainfiles = f.readlines()  # 标注数据 格式(123.jpg pig x_min y_min x_max y_max)

    file_names = []

    head = True

    folder = 'a'
    file_name = file_name
    path = filepath

    im = cv2.imread(filepath)  # 读取图片信息
    h = str(im.shape[0])
    w = str(im.shape[1])
    print(file_name)

    for line in trainfiles:
        trainFile = line.split()

        if head:
            head = False
            tree = objectify.ElementMaker(annotate=False)
            # 如果没有重复，则顺利进行。这给的数据集一张图片的多个框没有写在一起。
            lable = trainFile[0]
            xmin = str((float(trainFile[1])*im.shape[1]))
            ymin = str((float(trainFile[2])*im.shape[0]))
            xmax = str((float(trainFile[3])*im.shape[1]))
            ymax = str((float(trainFile[4])*im.shape[0]))
            print(xmin,ymin,xmax,ymax)

            anno_tree = tree.annotation(tree.folder(folder), tree.filename(file_name), tree.path(path),
                                        tree.source(tree.database('Unknown')),
                                        tree.size(tree.width(w), tree.height(h), tree.depth('3')),
                                        tree.segmented(0),
                                        tree.object(tree.name(lable), tree.pose('Unspecified'), tree.truncated(0),
                                                    tree.difficult(0),
                                                    tree.bndbox(tree.xmin(xmin), tree.ymin(ymin), tree.xmax(xmax),
                                                                tree.ymax(ymax))))

            etree.ElementTree(anno_tree).write(template_file, pretty_print=True)
        else:
            lable = trainFile[0]
            xmin = str((float(trainFile[1])))
            ymin = str((float(trainFile[2])))
            xmax = str((float(trainFile[3])))
            ymax = str((float(trainFile[4])))
            print(xmin,ymin,xmax,ymax)

            tree = etree.ElementTree()
            tree.parse(template_file)

            root = tree.getroot()
            element = etree.Element('object')

            name = etree.Element('name')
            name.text = lable
            element.append(name)

            pose = etree.Element('pose')
            pose.text = 'Unspecified'
            element.append(pose)

            truncated = etree.Element('truncated')
            truncated.text = '0'
            element.append(truncated)

            difficult = etree.Element('difficult')
            difficult.text = '0'
            element.append(difficult)

            lxmin = etree.Element('xmin')
            lxmin.text = xmin
            element.append(lxmin)

            lymin = etree.Element('ymin')
            lymin.text = ymin
            element.append(lymin)

            lxmax = etree.Element('xmax')
            lxmax.text = xmax
            element.append(lxmax)

            lymax = etree.Element('ymax')
            lymax.text = ymax
            element.append(lymax)

            root.append(element)

            tree.write(template_file, encoding='utf-8')


if __name__ == '__main__':
    txtfiles = os.listdir('ImageSets/v')
    xmldir = 'Annotations/b'

    for item in txtfiles:
        xmlpath = os.path.join('Annotations/b', item[:-3] + 'xml')
       
        txt2xml(template_file=xmlpath, file_name=item[:-3] + 'jpg', filepath=os.path.join('Annotations/a', item[:-3] + 'jpg'),
                train_file=os.path.join('ImageSets/v', item))
